from statistics import mode

import numpy as np
import cv2
import paho.mqtt.client as mqtt
import json
import time

from PIL import Image, ImageFont, ImageDraw
from rknn.api import RKNN
from wz_ai.inference import detect_faces
from wz_ai.inference import draw_text
from wz_ai.inference import draw_bounding_box
from wz_ai.inference import apply_offsets
from wz_ai.inference import load_detection_model
from wz_ai.preprocessor import preprocess_input

# parameters for loading data and images

print("XXXXXXXXXX   唯众AI实训平台  XXXXXXXXXX")
print(".........      平台启动中    ..........")
print(".........       请等待      ...........")

detection_model_path = './face_reg.xml'

def get_labels(dataset_name):
    if dataset_name == 'fer2013':
        return {0: 'angry', 1: 'disgust', 2: 'fear', 3: 'happy',
                4: 'sad', 5: 'surprise', 6: 'neutral'}

emotion_labels = get_labels('fer2013')

INPUT_SIZE = 64

def load_model(modle_path):
	# Create RKNN object
	rknn = RKNN()

	print("XXXXXXXXXXX      WZ_AI     XXXXXXXXXXX")
	print('-->开始加载模型')
	rknn.load_rknn(modle_path)
	print('模型加载成功')

	# init runtime environment
	print('--> 初始化环境中......')
	ret = rknn.init_runtime()
	if ret != 0:
		print('环境初始化失败')
		exit(ret)
	print('环境初始化完成')
	print("XXXXXXXXXXX      WZ_AI     XXXXXXXXXXX")
	return rknn

def model_predict(rknn, gray_face):
    emotion_prediction = rknn.inference(inputs=[gray_face], data_type='float32', data_format='nhwc')

    emotion_probability = np.max(emotion_prediction)
    emotion_label_arg = np.argmax(emotion_prediction)
    emotion_text = emotion_labels[emotion_label_arg]
    #print('emotion: ' + emotion_text)
    return emotion_text, emotion_probability


HOST = "139.9.246.47"
PORT = 18883



def on_message(client, userdata, message):
    print(message.topic + " " + ":" + str(message.payload))


def on_connect(client, userdata, flags, rc):
    print("Connected with result code " + str(rc))
    client.subscribe('kongzhi')


def client_pb(message:str):
    time_now = time.strftime('%Y-%m-%d %H-%M-%S', time.localtime(time.time()))
    payload = {'msg':'%s' % message, 'time':'%s' % time_now}
    client.publish('kongzhi', json.dumps(payload, ensure_ascii=False))
    print('publish success!!')
    return True
    
if __name__ == '__main__':
    # 获取模型输入尺寸
    emotion_target_size = (INPUT_SIZE, INPUT_SIZE)
    # hyper-parameters for bounding boxes shape
    frame_window = 10
    emotion_offsets = (20, 40)

    # starting lists for calculating modes
    emotion_window = []

    # 加载模型
    rknn = load_model('./biaoqing.rknn')
    face_detection = load_detection_model(detection_model_path)

    # starting video streaming
    cv2.namedWindow('WZ_AI_Platform ( Please press Esc to exit !)', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('WZ_AI_Platform ( Please press Esc to exit !)', 1920, 1080)
    video_capture = cv2.VideoCapture(0)
    
    client = mqtt.Client(client_id='whwzzc_bq')
    client.connect('139.9.246.47', 18883, 60)
    client.username_pw_set('admin', 'public')
    client.loop_start()
   
    emotion_text = ''
    emotion_old = 'nothing'
    while True:
        frame_start = time.time()
        bgr_image = video_capture.read()[1]
        gray_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
        rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
        faces = detect_faces(face_detection, gray_image)
        print('人脸识别耗时: %f'%(time.time() - frame_start)+' s')

        for face_coordinates in faces:
            start = time.time()
            x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
            gray_face = gray_image[y1:y2, x1:x2]
            try:
                gray_face = cv2.resize(gray_face, (emotion_target_size))
            except:
                continue

            gray_face = preprocess_input(gray_face, True)
            gray_face = np.expand_dims(gray_face, 0)
            gray_face = np.expand_dims(gray_face, -1)
            emotion_text, emotion_probability = model_predict(rknn, gray_face)
            print('表情识别耗时: %f'%(time.time() - start)+' s')
            emotion_window.append(emotion_text)

            if len(emotion_window) > frame_window:
                emotion_window.pop(0)
            try:
                emotion_mode = mode(emotion_window)
            except:
                continue

            if emotion_text == 'angry':
                color = emotion_probability * np.asarray((255, 0, 0))
            elif emotion_text == 'sad':
                color = emotion_probability * np.asarray((0, 0, 255))
            elif emotion_text == 'happy':
                color = emotion_probability * np.asarray((255, 255, 0))
            elif emotion_text == 'surprise':
                color = emotion_probability * np.asarray((0, 255, 255))
            else:
                color = emotion_probability * np.asarray((0, 255, 0))
            color = color.astype(int)
            color = color.tolist()
            draw_bounding_box(face_coordinates, rgb_image, color)
            draw_text(face_coordinates, rgb_image, emotion_mode,
                      color, 0, 0, 1, 1)

        fps_text = ("%.2f" % (1 / (time.time() - frame_start)))
        cv2.putText(rgb_image, 'FPS: ' +fps_text, (10, 20),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5, (0, 255, 0), 1, cv2.LINE_AA)


        bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)

        change_img = Image.fromarray(rgb_image)

        # PIL图片上打印汉字
        draw = ImageDraw.Draw(change_img)  # 图片上打印
        font = ImageFont.truetype("wz_ziti.ttf", 30, encoding="utf-8")  # 参数1：字体文件路径，参数2：字体大小
        draw.text((220, 0), "唯众智能灯控系统", (255, 255, 0), font=font)  # 参数1：打印坐标，参数2：文本，参数3：字体颜色，参数4：字体
        font = ImageFont.truetype("wz_ziti.ttf", 20, encoding="utf-8")  
        draw.text((215, 40), "(根据表情改变灯光的颜色)", (255, 0, 0), font=font)  
        cv2charimg = cv2.cvtColor(np.array(change_img), cv2.COLOR_RGB2BGR)


        cv2.imshow('WZ_AI_Platform ( Please press Esc to exit !)', cv2charimg)

        # cv2.imshow('WZ_AI_Platform ( Please press Esc to exit !)', bgr_image)

       
        if (emotion_text != '') and (emotion_old != emotion_text):
            emotion_old = emotion_text
            client_pb(emotion_text)
            # publish.single('kongzhi', emotion_text, qos=1, hostname=HOST, port=PORT, client_id='wzai',auth={'username':'admin', 'password':'public'})
        

        if cv2.waitKey(1) & 0xFF == 27:
            break
